#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : RAG-demo
# @File    : elasticsearch_util.py
# @IDE     : PyCharm
# @Author  :ZH
# @Time    : 2025/6/27 17:40
from elasticsearch import Elasticsearch, NotFoundError

from util.logger_setup_util import LoggerSetup

logger = LoggerSetup(__name__).get_logger()


class ElasticsearchBM25Client:
    def __init__(self, host="localhost", port=9200, index_name="bm25-index"):
        self.es = Elasticsearch([f"http://{host}:{port}"])
        self.index_name = index_name

    def create_index(self):
        """创建索引，包含 text 和 filename 两个字段"""
        body = {
            "settings": {
                "analysis": {
                    "analyzer": {
                        "default": {
                            "type": "standard"
                        }
                    }
                },
                "index": {
                    "similarity": {
                        "default": {
                            "type": "BM25"
                        }
                    }
                }
            },
            "mappings": {
                "properties": {
                    "content": {"type": "text"},
                    "filename": {"type": "keyword"},  # 用于存储文件名
                    "knowledge_id": {"type": "keyword"},  # 用于存储知识库ID
                    "file_hase": {"type": "keyword"}  # 用于存储文件哈希值
                }
            }
        }
        if not self.es.indices.exists(index=self.index_name):
            self.es.indices.create(index=self.index_name, body=body)
            logger.info(f"Index '{self.index_name}' created.")
        else:
            logger.info(f"Index '{self.index_name}' already exists.")

    def delete_index(self):
        """删除指定索引"""
        if self.es.indices.exists(index=self.index_name):
            self.es.indices.delete(index=self.index_name)
            logger.info(f"Index '{self.index_name}' deleted.")
        else:
            logger.info(f"Index '{self.index_name}' does not exist.")

    def add_document(self, doc_id, text, filename):
        """添加文档，包含文本内容和文件名"""
        body = {
            "content": text,
            "filename": filename,
            "knowledge_id": None,  # 可选字段，默认为 None
            "file_hase": None  # 可选字段，默认为 None
        }
        self.es.index(index=self.index_name, id=doc_id, document=body)

    def add_documents(self, documents):
        """批量添加文档"""
        body = []
        for doc in documents:
            meta = {"index": {"_index": self.index_name}}
            data = {
                "content": doc["document"],
                "filename": doc.get("file_name", None),
                "knowledge_id": doc.get("knowledge_id", None),
                "file_hase": doc.get("file_hase", None)
            }
            body.append(meta)
            body.append(data)

        if body:
            self.es.bulk(body=body)
            logger.info(f"Added {len(documents)} documents to index '{self.index_name}'.")

    def delete_document(self, knowledge_id, file_hase: str = None):
        """删除文档"""
        try:
            # 删除文档
            self.es.delete_by_query(
                index=self.index_name,
                body={
                    "query": {
                        "bool": {
                            "must": [
                                {"match": {"knowledge_id": knowledge_id}},
                                {"match": {"file_hase": file_hase}}
                            ] if file_hase else [
                                {"match": {"knowledge_id": knowledge_id}}
                            ]
                        }
                    }
                }
            )
            logger.info(f"Document with knowledge_id '{knowledge_id}' and file_hase '{file_hase}' deleted from index '{self.index_name}'.")
        except NotFoundError:
            logger.error(f"Document with knowledge_id '{knowledge_id}' and file_hase '{file_hase}' not found in index '{self.index_name}'.")

    def keyword_search(self, query_text, size=10):
        """基于关键词的 BM25 搜索，返回结果包括 text 和 filename"""
        body = {
            "query": {
                "match": {
                    "text": query_text
                }
            },
            "size": size
        }
        result = self.es.search(index=self.index_name, body=body)
        return [
            (hit["_id"], hit["_score"], hit["_source"]["text"], hit["_source"].get("filename", None))
            for hit in result["hits"]["hits"]
        ]
